package com.tkk.monitoring.analyze.exp.sql;

import com.tkk.monitoring.analyze.db.DB;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.log4j.Logger;

import java.io.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * author: Tkk
 * date: 2015/7/8
 */
public class SQLMapWorker implements Runnable {

    public static final BlockingQueue<SQLMapWork> QUEUE = new LinkedBlockingQueue<SQLMapWork>(100);
    public static final Executor executor = Executors.newFixedThreadPool(10);
    public static final Logger logger = Logger.getLogger(SQLMapWorker.class);

    static {
        for (int i = 0; i < 10; i++) {
            executor.execute(new SQLMapWorker());
        }
    }

    @Override
    public void run() {
        while (true) {
            PrintWriter wr = null;
            BufferedReader br = null;
            try {
                SQLMapWork take = QUEUE.take();

                //
                File logFile = new File(take.getLogFilePath());
                String command = StringUtils.join(take.getCommand(), " ");
                logger.info(String.format("SQLMAP[开始]: %s", command));

                //
                Process process = Runtime.getRuntime().exec(command);
                wr = new PrintWriter(new FileOutputStream(logFile));
                wr.println("-----------------------------");
                wr.println(command);
                wr.println("-----------------------------");
                wr.flush();
                br = new BufferedReader(new InputStreamReader(process.getInputStream()));
                String output;
                while (null != (output = br.readLine())) {
                    wr.println(output);
                    wr.flush();
                }

                //
                process.waitFor();
                logger.info(String.format("SQLMAP[结束]: %s", command));

                //
                String logContent = FileUtils.readFileToString(logFile);
                if (logContent != null && (logContent.indexOf("is vulnerable") != -1 || logContent.indexOf("Parameter:") != -1)) {
                    logger.warn(String.format("SQLMAP[发现]: %s", command));
                    DB.getInstance().addPoc("SQL注入", take.getTarget().getUrl(), logContent, "");
                }
            } catch (InterruptedException e) {
                e.printStackTrace();
            } catch (IOException e) {
                e.printStackTrace();
            } finally {
                if (wr != null) {
                    wr.close();
                }
                if (br != null) {
                    try {
                        br.close();
                    } catch (IOException e) {
                    }
                }
            }
        }
    }
}
